"""
基于 AutoGen 的投资组合优化流程
"""
from agents.stock_agents import StockAgentFactory
from autogen import GroupChat, GroupChatManager

class PortfolioOptimizationFlow:
    def __init__(self, stock_list):
        self.stock_list = stock_list
        self.agent_factory = StockAgentFactory()
        self.agents = self.agent_factory.create_all_agents(stock_list[0]) #  For now, just use the first stock to initialize

    def run(self):
        # 1. Market Research
        market_research_report = self.agents["market_researcher"].initiate_chat(
            self.agents["user_proxy"], message=f"分析当前市场趋势，尤其是与 {self.stock_list} 相关的行业。"
        )

        # 2. Risk a
        risk_analysis_reports = {}
        for stock in self.stock_list:
            risk_analysis_reports[stock] = self.agents["risk_manager"].initiate_chat(
                self.agents["user_proxy"], message=f"评估 {stock} 的风险。"
            )

        # 3. Portfolio Construction
        portfolio_construction_report = self.agents["portfolio_manager"].initiate_chat(
            self.agents["user_proxy"],
            message=f"""
            根据以下信息构建一个投资组合：
            市场研究报告: {market_research_report}
            风险分析报告: {risk_analysis_reports}
            投资组合: {self.stock_list}
            """,
        )

        return portfolio_construction_report

